from torch.utils.data import Dataset,DataLoader
import torchvision.datasets as datasets
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from torch.optim import SGD

data_folder = './FMNIST'
fmnist = datasets.FashionMNIST(data_folder,download=True,train=True)
device = 'cpu'

tr_images = fmnist.data
tr_targets = fmnist.targets

class FMNISTDataset(Dataset):
    def __init__(self,x,y):
        x = x.float() / 255
        self.x = x.view(-1,784)
        self.y = y

    def __getitem__(self, item):
        return self.x[item],self.y[item]

    def __len__(self):
        return len(self.x)


def getdata():
    train = FMNISTDataset(tr_images,tr_targets)
    trn_dl = DataLoader(train,batch_size=32,shuffle=True)
    return trn_dl

def get_model():
    model = nn.Sequential(
        nn.Linear(784,1000),
        nn.ReLU(),
        nn.Linear(1000,10)
    )
    loss_fn = nn.CrossEntropyLoss()
    optimizer = SGD(model.parameters(),lr=0.000001)
    return model,loss_fn,optimizer

def train_batch(x,y,model,opt,loss_fn):
    model.train()
    prediction = model(x)
    loss_value = loss_fn(prediction,y)
    loss_value.backward()
    opt.step()
    opt.zero_grad()
    return loss_value.item()

@torch.no_grad()
def accuracy(x,y,model):
    model.eval()
    prediction = model(x)
    max_values,argmaxes = prediction.max(-1)
    is_correct = argmaxes == y
    return is_correct.cpu().numpy().tolist()

trn_dl = getdata()
model,loss_fn,optimizer = get_model()
losses , accuracies = [],[]
eps = 100
for epoch in range(eps):
    print(epoch)
    epoch_losses , epoch_accuracies = [],[]
    for ix, batch in enumerate(iter(trn_dl)):
        x, y = batch
        batch_loss = train_batch(x, y, model, optimizer, loss_fn)
        epoch_losses.append(batch_loss)
    epoch_loss = np.array(epoch_losses).mean()
    for ix, batch in enumerate(iter(trn_dl)):
        x, y = batch
        is_correct = accuracy(x, y, model)
        epoch_accuracies.extend(is_correct)
    epoch_accuracy = np.mean(epoch_accuracies)
    losses.append(epoch_loss)
    accuracies.append(epoch_accuracy)

epochs = np.arange(eps) + 1
plt.figure(figsize=(20, 5))
plt.subplot(121)
plt.title('Loss value over increasing epochs')
plt.plot(epochs, losses, label='Training Loss')
plt.legend()
plt.subplot(122)
plt.title('Accuracy value over increasing epochs')
plt.plot(epochs, accuracies, label='Training Accuracy')
plt.gca().set_yticklabels(['{:.0f}%'.format(x*100) for x in plt.gca().get_yticks()])
plt.legend()
plt.show()
